#include <iostream>
#include <cstdlib>
#include <string>
#include <time.h>
#include <thread>
#include <queue>
#include <unistd.h>

#include "opencv2/opencv.hpp"
#include <opencv2/core/core.hpp>
#include <opencv2/highgui/highgui.hpp>  
#include <opencv2/imgproc/imgproc.hpp>  
#include <opencv2/highgui/highgui_c.h>

#include "nanotrack.hpp" 
#include "camera.hpp"

using namespace std;
using namespace cv;

std::queue<cv::Mat> readQueue, dispQueue;
cv::Point selectedPoint; 
cv::Rect trackWindow;
cv::Mat frame;

bool trackEnabled = false; // 控制是否开启tracking
bool mouseSelectRect = false; // 控制选择模式，true为框选，false为点选
bool updateTracker = false; // 控制是否重新初始化tracking对象
int searchRange = 150;

void onMouse(int event, int x, int y, int, void* userdata)  
{  
    if (event == cv::EVENT_LBUTTONDOWN)  
    {  
        cout << "Mouse click at: " << x << ", " << y << endl;
        if (mouseSelectRect) {
            // 如果是框选模式，则需要设置初始的trackWindow
            trackWindow = cv::Rect(1000, 1000, 100, 100);
            // trackWindow = cv::selectROI("window", frame); 
            std::cout << "Select a rect done." << std::endl;
        } else {
            // 如果是点选模式，则只记录点的位置
            // selectedPoint = cv::Point(x, y);
            // std::cout << "Select a point done." << std::endl;
            
            cv::Mat* img = reinterpret_cast<cv::Mat*>(userdata);
            int minX = std::max(x - searchRange, 0);
            int maxX = std::min(x + searchRange, img->cols - 1);
            int minY = std::max(y - searchRange, 0);
            int maxY = std::min(y + searchRange, img->rows - 1);
            trackWindow = cv::Rect(minX, minY, maxX - minX, maxY - minY);
            
            for (int i = minY; i <= maxY; i++) {
                for (int j = minX; j <= maxX; j++) {
                    cv::Vec3b color = img->at<cv::Vec3b>(i, j);
                    if (color[0] < 127 && color[1] < 127 && color[2] < 127) { // Check if the pixel is non-white
                        trackWindow.x = std::min(trackWindow.x, j);
                        trackWindow.y = std::min(trackWindow.y, i);
                        trackWindow.width = std::max(trackWindow.width, j - trackWindow.x);
                        trackWindow.height = std::max(trackWindow.height, i - trackWindow.y);
                    }
                }
            }
        }
        trackEnabled = true;
        updateTracker = true; // 需要更新跟踪器
    }  
} 

void cxy_wh_2_rect(const cv::Point& pos, const cv::Point2f& sz, cv::Rect &rect)
{
    rect.x = cv::max(0, pos.x - int(sz.x / 2));
    rect.y = cv::max(0, pos.y - int(sz.y / 2));
    rect.width = int(sz.x);   
    rect.height = int(sz.y);  
    
}

int main(int argc, char** argv)
{
    std::string T_backbone_model = "../models/track_backbone_T.rknn";
    std::string X_backbone_model = "../models/track_backbone_X.rknn";
    std::string head_model = "../models/head.rknn";
    const char* video_path = "../data/videos/visDrone.mp4"; 

    // Read video
    CameraVideoCapture capture(1920, 1080, "/dev/video8", 30);
    frame = capture.getFrame();

    //display
    cv::namedWindow("window", CV_WINDOW_NORMAL);
    cv::setWindowProperty("window", CV_WND_PROP_FULLSCREEN, CV_WINDOW_FULLSCREEN);

    NanoTrack *siam_tracker; 
    siam_tracker = new NanoTrack(T_backbone_model, X_backbone_model, head_model); 
    clock_t start = clock();

    State state; 
    cv::setMouseCallback("window", onMouse, reinterpret_cast<void*>(&frame));  
    float cls_score_max = 0;
    
    while (true){
        frame = capture.getFrame();

        // 更新跟踪目标
        if (updateTracker){
            cout << "Init trackWindow using " << (mouseSelectRect ? "SelectRect" : "SelectPoint") << std::endl;
            siam_tracker->init(frame, trackWindow); 
            updateTracker = false;
            std::cout << "==========================" << std::endl;
            std::cout << "Re-init done!" << std::endl << std::endl; 
        }
        
        // 绘制跟踪框
        if (trackEnabled) {
            cls_score_max = siam_tracker->track(frame); 
            cv::Rect rect;
            cxy_wh_2_rect(siam_tracker->state.target_pos, siam_tracker->state.target_sz, rect);
            cv::putText(frame,to_string(cls_score_max),cv::Point(rect.x,rect.y),cv::FONT_HERSHEY_SIMPLEX,2,cv::Scalar(0,0,255),2);
            cv::rectangle(frame, rect, cv::Scalar(0, 255, 0), 3);
            
        }

        // std::cout << "fps : " << CLOCKS_PER_SEC / double(clock() - start) << std::endl;
        start = clock();

        cv::imshow("window", frame);
        
        // 键盘切换模式（s：更改选择方式，r：是否显示跟踪效果）
        int key = cv::waitKey(1);
        if (static_cast<char>(key) == 's') {
            mouseSelectRect = !mouseSelectRect;
            trackEnabled = !trackEnabled;
            std::cout << "Pressed 's' key." << std::endl;
            if (mouseSelectRect) std::cout << "Change mode to select rect." << std::endl << std::endl;
            else std::cout << "Change mode to select point." << std::endl << std::endl;
        }
        if (static_cast<char>(key) == 'r') {
            trackEnabled = !trackEnabled;
            cout << "Whether display: " << trackEnabled << std::endl;
        }
    }
    cv::destroyWindow("window");
    return 0;
}
